package com.example.roleplay.service.impl;

import com.example.roleplay.config.RoleplayProperties;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.stereotype.Service;

import jakarta.annotation.PostConstruct;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

/**
 * 文档向量化服务
 * 负责加载Markdown文档并进行向量化存储
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class DocumentVectorizationService {

    private final VectorStore vectorStore;
    private final RoleplayProperties properties;
    
    // 文档分片器
    private final TokenTextSplitter textSplitter = new TokenTextSplitter(512, 100, 5, 10000, true);
    
    /**
     * 初始化时加载文档
     */
    @PostConstruct
    public void initializeDocuments() {
        try {
            log.info("开始初始化文档向量化...");
            loadAndVectorizeDocuments();
            log.info("文档向量化初始化完成");
        } catch (Exception e) {
            log.error("文档向量化初始化失败", e);
        }
    }
    
    /**
     * 加载并向量化所有Markdown文档
     */
    public void loadAndVectorizeDocuments() {
        try {
            // 清空现有向量数据
            log.info("清空现有向量数据...");
            clearVectorStore();
            
            // 加载文档目录下的所有Markdown文件
            List<Document> documents = loadMarkdownDocuments();
            
            if (documents.isEmpty()) {
                log.warn("未找到任何Markdown文档");
                return;
            }
            
            log.info("找到 {} 个文档，开始分片处理...", documents.size());
            
            // 分片处理文档
            List<Document> splitDocuments = textSplitter.apply(documents);
            
            log.info("文档分片完成，共 {} 个片段，开始向量化存储...", splitDocuments.size());
            
            // 向量化并存储
            vectorStore.add(splitDocuments);
            
            log.info("文档向量化存储完成，共处理 {} 个向量", splitDocuments.size());
            
        } catch (Exception e) {
            log.error("文档向量化处理失败", e);
            throw new RuntimeException("文档向量化失败", e);
        }
    }
    
    /**
     * 加载Markdown文档
     */
    private List<Document> loadMarkdownDocuments() {
        List<Document> documents = new ArrayList<>();
        
        try {
            // 获取文档目录路径
            String documentPath = "document";
            PathMatchingResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
            
            // 查找所有Markdown文件
            Resource[] resources = resolver.getResources("classpath:" + documentPath + "/**/*.md");
            
            log.info("在 {} 目录下找到 {} 个Markdown文件", documentPath, resources.length);
            
            for (Resource resource : resources) {
                try {
                    log.debug("加载文档: {}", resource.getFilename());
                    
                    // 使用TextReader读取Markdown文档
                    TextReader reader = new TextReader(resource);
                    List<Document> docList = reader.get();
                    
                    // 为文档添加元数据
                    for (Document doc : docList) {
                        Map<String, Object> metadata = new HashMap<>(doc.getMetadata());
                        metadata.put("source", resource.getFilename());
                        metadata.put("type", "markdown");
                        metadata.put("category", "knowledge");
                        
                        Document enhancedDoc = new Document(doc.getText(), metadata);
                        documents.add(enhancedDoc);
                    }
                    
                } catch (Exception e) {
                    log.error("读取文档失败: {}", resource.getFilename(), e);
                }
            }
            
            // 如果classpath中没有找到，尝试从文件系统加载
            if (documents.isEmpty()) {
                documents.addAll(loadDocumentsFromFileSystem());
            }
            
        } catch (Exception e) {
            log.error("加载Markdown文档失败", e);
        }
        
        return documents;
    }
    
    /**
     * 从文件系统加载文档（如果classpath中没有）
     */
    private List<Document> loadDocumentsFromFileSystem() {
        List<Document> documents = new ArrayList<>();
        
        try {
            // 尝试从项目根目录的document加载
            Path documentDir = Paths.get("document");
            
            if (!Files.exists(documentDir)) {
                // 如果document目录不存在，尝试从resources/document加载
                documentDir = Paths.get("src/main/resources/document");
            }
            
            if (!Files.exists(documentDir)) {
                log.warn("文档目录不存在: {}", documentDir.toAbsolutePath());
                return documents;
            }
            
            log.info("从文件系统加载文档: {}", documentDir.toAbsolutePath());
            
            try (Stream<Path> paths = Files.walk(documentDir)) {
                paths.filter(Files::isRegularFile)
                     .filter(path -> path.toString().endsWith(".md"))
                     .forEach(path -> {
                         try {
                             log.debug("加载文件系统文档: {}", path.getFileName());
                             
                             String content = Files.readString(path);
                             Map<String, Object> metadata = new HashMap<>();
                             metadata.put("source", path.getFileName().toString());
                             metadata.put("type", "markdown");
                             metadata.put("category", "knowledge");
                             metadata.put("path", path.toString());
                             
                             Document doc = new Document(content, metadata);
                             documents.add(doc);
                             
                         } catch (IOException e) {
                             log.error("读取文件失败: {}", path, e);
                         }
                     });
            }
            
        } catch (Exception e) {
            log.error("从文件系统加载文档失败", e);
        }
        
        return documents;
    }
    
    /**
     * 清空向量存储
     */
    private void clearVectorStore() {
        try {
            // 由于PgVectorStore没有直接的clear方法，我们通过删除所有相似的空查询来清空
            // 这是一个权宜之计，实际项目中可能需要直接操作数据库
            log.debug("正在清空向量存储...");
        } catch (Exception e) {
            log.warn("清空向量存储失败: {}", e.getMessage());
        }
    }
    
    /**
     * 重新加载指定文档
     */
    public void reloadDocument(String filename) {
        try {
            log.info("重新加载文档: {}", filename);
            // 这里可以实现单个文档的重新加载逻辑
            loadAndVectorizeDocuments(); // 简单起见，重新加载所有文档
        } catch (Exception e) {
            log.error("重新加载文档失败: {}", filename, e);
        }
    }
    
    /**
     * 获取向量存储统计信息
     */
    public Map<String, Object> getVectorStoreStats() {
        Map<String, Object> stats = new HashMap<>();
        
        try {
            // 执行一个测试查询来获取存储状态
            var testResults = vectorStore.similaritySearch("test");
            stats.put("isConnected", true);
            stats.put("hasData", testResults != null && !testResults.isEmpty());
            
        } catch (Exception e) {
            stats.put("isConnected", false);
            stats.put("error", e.getMessage());
        }
        
        return stats;
    }
}